'''视觉处理模块'''

import cv2
import torch
from mmdet.apis import inference_detector, init_detector
from configs import *
import copy,time,os,sys
import numpy as np
dist_gap = 15#像素
import threading
from module_order import HttpOrder
from module_robot import HandRobot
from module_vision_similar import *
from hkcamera_reader import *
import threading
class StereoVision:
    def __init__(self, frame_interval: int = 0.03):
        self.last_result = {"labels":[],"scores":[],"bbox":[]}
        self.current_result={"labels":[],"scores":[],"bbox":[]}
        self.current_result_hk = {"labels":[],"scores":[],"bbox":[]}
        self.ready=False
        self.ready_hk = False
        self.daocai_ready = False
        self.img_draw_finish = False
        self.img_draw_finish_hk = False
        self.img_hk_draw_finish = False
        self.args = parse_args()
        self.edge = np.array(self.args.max_vision_edge)
        self.edge_hk = np.array(self.args.max_vision_edge_hk)
        self.edge_flag = np.array([True,True,False,False])#坐标差满足该形式，代表在区域内
        self.plate_exisit_flag = False
        self.plate_msg = {}
        self.pot_msg = {}#TODO 目前不能放同样的菜
        self.pot_plate_msg_done = False
        self.stander_imgs_pot = []
        self.stander_imgs_plate = []
        self.init_stander_imgs(self.args.stander_plate_dir,self.args.stander_pot_dir)
        self.plat_pot_clock = threading.Lock()
        # self.init_stander_imgs()
        # self.handrobot = HandRobot(self.args.robot_ip)
        # self.flow=[False,False,False]
        # self.stereovision = StereoVision()

        # frame parameters
        self.frame_interval = frame_interval

        self.img = None
        self.img_hk = None
        self.depth_img=None
        self.time_last = time.time()

        self.load_model()
        self.hk_reader = HKCamera()
        
    def init_stander_imgs(self,img_plate_dir,img_pot_dir):
        self.faisssearch = FaissSearch()
        self.stander_imgs_plate = os.listdir(img_plate_dir)
        # for img in stander_imgs_name:
        #     stander_img = cv2.imread(os.path.join(img_plate_dir, img))
        #     self.stander_imgs_plate.append((img.split('.')[0],stander_img))

        self.stander_imgs_pot = os.listdir(img_pot_dir)
        # for img in stander_imgs_name2:
        #     stander_img = cv2.imread(os.path.join(img_pot_dir, img))
        #     self.stander_imgs_pot.append((img.split('.')[0],stander_img))
    def pixel2world(self,pixel_coord:tuple):
        #像素点转世界点
        robot_home = (650,492)
        offset = (180,10)#世界坐标系，单位mm
        #92 1801 --> 1500mm -->1709
        #285 1015 -->1500mm -->730
        ratio = 0.48#1.23
        x,y = pixel_coord
        x = x-robot_home[0]#坐标原点移到home点
        y = y-robot_home[1]
        #像素坐标xy与世界坐标xy交换
        world_x=int(y/ratio)
        world_y = int(x/ratio)
        return [world_x+offset[0],world_y+offset[1]]
    
    def x_pixel2world(self,point_x):
        return self.pixel2world((point_x,0))[0]

    def y_pixel2world(self,point_y):
        return self.pixel2world((0,point_y))[1]
    
    def get_area_depth(self,pixel_coord):
        #获取指定点附近的物体平均深度
        depth = get_depth_at_point(self.depth_img,pixel_coord[0],pixel_coord[1])
        return depth-1650#摄像头距离盆地1650mm

    def load_model(self):
        device = torch.device(self.args.device)
        self.model = init_detector(self.args.config, self.args.checkpoint, device=device)

    def check_daocai_ready(self):#判断餐盘是否符合倒菜条件--一定时间内餐盘两个监测点的坐标基本不变
        # {'scores': array([0.9507748 , 0.94361126, 0.93571883, 0.8750038 , 0.86429864,
        #     0.79122   , 0.7179362 ], dtype=float32), 'bbox': array([[ 717.61694,  205.31859,  843.8498 ,  333.5363 ],
        #     [ 855.3349 ,  216.56024,  987.74475,  345.32333],
        #     [ 704.92035,  352.61554,  833.2615 ,  481.9933 ],
        #     [ 638.61487,  194.64867, 1150.12   ,  525.52374],
        #     [ 979.343  ,  373.5901 , 1122.9742 ,  518.5199 ],
        #     [ 998.7687 ,  224.87195, 1130.9976 ,  360.86182],
        #     [ 997.9946 ,  226.1797 , 1129.7852 ,  359.95154]], dtype=float32), 'labels': array([1, 2, 4, 0, 4, 1, 3], dtype=int64)}
        for lable,bbox in zip(self.last_result["labels"],self.last_result["bbox"]):
            if lable =="0":
                # last_result = bbox
                last_result = bbox
                break
        else:
            last_result =np.array([])
        for lable,bbox in zip(self.current_result["labels"],self.current_result["bbox"]):
            if lable=="0":
                current_result = bbox
                break
        else:
            current_result =np.array([])
        if last_result.size and current_result.size:
            abs_result = abs(current_result-last_result)
            four_points_remain = [(np.array([abs_result[i],abs_result[i+1]])<dist_gap).all() for i in range(0,len(abs_result),2)]#四个点是否变化
            # print("餐盘角点变化情况{}".format(four_points_remain))
            if np.array(four_points_remain).sum()>2:#有三四个点没有变化
                self.daocai_ready = True
            else :
                # self.time_last = time.time()
                self.daocai_ready = False
        else:
            # self.time_last = time.time()
            self.daocai_ready = False
        # print(self.daocai_ready)

    def caculate_similar(self,input_img,stander_imgs,flag):
        #计算相似度，返回改img对应的内容
        if flag =="plate":
            dis,index = self.faisssearch.search(input_img,"vector_plate.index")
            # print("similarplate",dis)
            if dis[0][0]<1.2:
                return stander_imgs[index[0][0]].split(".")[0]
            else:
                return "1"
        else:
            dis,index = self.faisssearch.search(input_img,"vector_pot.index")
            # print("similarpot",dis)
            # if dis[0][0]<0.6:
            #     return stander_imgs[index[0][0]].split(".")[0]
            # else:
            #     return "0"
            return stander_imgs[index[0][0]].split(".")[0]

    def change_label(self,labels:np.array,bboxes:np.array):
        #根据相对位置，计算出目标的label
        whitch_side = True
        labels = [str(l) for l in labels.tolist()]
        plate_bboxe,plate_chopstick_bboxe,plate_circle_bboxe= np.array([]),np.array([]),np.array([])
        for label,bboxe in zip(labels,bboxes):
            if label == "0":#plate
                plate_bboxe = bboxe
                self.plate_exisit_flag = True
                break
        else:
            self.plate_exisit_flag=False
        if plate_bboxe.size:
            for label,bboxe in zip(labels,bboxes):
                if label == "1":#plate_chopsticks
                    # print(bboxe,plate_bboxe)
                    if abs(bboxe[0]-plate_bboxe[0])>200:
                        whitch_side = True
                    else:
                        whitch_side = False
                if label == "3":#plate_circle
                    plate_circle_bboxe = bboxe
        
        self.plat_pot_clock.acquire()
        self.pot_msg={}
        self.plate_msg={}
        self.pot_plate_msg_done=False
        indices = [i for i in range(len(bboxes))]
        for indice,label,bboxe in zip(indices,labels,bboxes):
            class_name = self.caculate_similar(self.img[int(bboxe[1]):int(bboxe[3]),int(bboxe[0]):int(bboxe[2])],self.stander_imgs_plate,flag = "plate")
            if plate_bboxe.size:
                if whitch_side:
                    if label == "2":#plate_blank
                        abs_x = abs(bboxe[0]-plate_bboxe[0])
                        abs_y = abs(bboxe[1]-plate_bboxe[1])
                        if abs_x>0 and abs_x<60:
                            labels[indice] = "p_11_"+class_name
                            self.plate_msg[plate_pos[0]]=class_name
                        elif abs_x>=60 and abs_x<180:
                            if abs_y>0 and abs_y<80:
                                labels[indice] ="p_12_"+class_name
                                self.plate_msg[plate_pos[1]]=class_name
                            else:
                                labels[indice] ="p_22_"+class_name
                                self.plate_msg[plate_pos[4]]=class_name
                        elif abs_x>=180 and abs_x<300:
                            labels[indice] ="p_13_"+class_name
                            self.plate_msg[plate_pos[2]]=class_name
                        else:
                            labels[indice] = "p_999"
                    if label =="3":
                        labels[indice] = "p_21_"+class_name
                        self.plate_msg[plate_pos[3]]="1"
                else:
                    if label == "2":#p_blank
                        abs_x = abs(bboxe[0]-plate_bboxe[0])
                        abs_y = abs(bboxe[1]-plate_bboxe[1])
                        if abs_x>0 and abs_x<130:
                            if abs_y>0 and abs_y<90:
                                labels[indice] = "p_22_"+class_name
                                self.plate_msg[plate_pos[4]]=class_name
                            else:
                                labels[indice] = "p_13_"+class_name
                                self.plate_msg[plate_pos[2]]=class_name
                        elif abs_x>=90 and abs_x<240:
                            labels[indice] = "p_12_"+class_name
                            self.plate_msg[plate_pos[1]]=class_name
                        elif abs_x>=240 and abs_x<360:
                            labels[indice] = "p_11_"+class_name
                            self.plate_msg[plate_pos[0]]=class_name
                        else:
                            labels[indice] ="p_999"
                    if label =="3":
                        labels[indice] = "p_21_"+class_name
                        self.plate_msg[plate_pos[3]]=class_name
            if label == "4":
                #TODO 相似度比较当前位置的菜品，获取菜品编号,label 改为  位置_编号
                abs_x = abs(bboxe[0])
                abs_y = abs(bboxe[1])
                class_name_pot = self.caculate_similar(self.img[int(bboxe[1]):int(bboxe[3]),int(bboxe[0]):int(bboxe[2])],self.stander_imgs_pot,flag = "pot")
                if abs_x>0 and abs_x<500:
                    if abs_y>0 and abs_y<350:
                        labels[indice] = "pot_11_"+class_name_pot
                        self.pot_msg[pot_pos[0]]=class_name_pot
                    elif abs_y>=350 and abs_y<750:
                        labels[indice] = "pot_21_"+class_name_pot
                        self.pot_msg[pot_pos[2]]=class_name_pot
                    elif abs_y>=750 and abs_x<1100:
                        labels[indice] = "pot_31_"+class_name_pot
                        self.pot_msg[pot_pos[4]]=class_name_pot
                elif abs_x>=500 and abs_x<1500:
                    if abs_y>0 and abs_y<350:
                        labels[indice] = "pot_12_"+class_name_pot
                        self.pot_msg[pot_pos[1]]=class_name_pot
                    elif abs_y>=350 and abs_y<750:
                        labels[indice] = "pot_22_"+class_name_pot
                        self.pot_msg[pot_pos[3]]=class_name_pot
                    elif abs_y>=750 and abs_y<1100:
                        labels[indice] = "pot_32_"+class_name_pot
                        self.pot_msg[pot_pos[5]]=class_name_pot
            # if bool(self.plate_msg):
            #     for k in sorted(self.plate_msg.keys()):
            #         self.
        # self.pot_msg={"11":"MDJ777","12":"MDJ867","21":"MDJ265","22":"MDJ001"}
        self.plat_pot_clock.release()
        
        return labels,bboxes
    
    def return_vision_result_hk(self):
        # camera = cv2.VideoCapture(0,cv2.CAP_MSMF)
        # camera.set(3, 1280)
        # camera.set(4, 1440)
        # cv2.namedWindow('camera',cv2.WINDOW_NORMAL)
        # time_current = time.time()
        # ch:开始取流 | en:Start grabbing
        ret=self.hk_reader.camera.MV3D_RGBD_Start()
        if ret != 0:
            print ("start fail! ret[0x%x]" % ret)
            self.hk_reader.camera.MV3D_RGBD_CloseDevice()
            os.system('pause')
            sys.exit()

        stFrameData = MV3D_RGBD_FRAME_DATA()  # OPT: move outside of while loop
        while True:
            time.sleep(self.frame_interval)    # OPT: change to 0.03, fps: 30

            ret=self.hk_reader.camera.MV3D_RGBD_FetchFrame(pointer(stFrameData), 5000)
            if ret==0:
                self.img_hk,self.depth_img = parse_frame(stFrameData)
                try:
                    # if time_current-self.time_last>time_gap:#记录time_gap之前的结果
                    #     self.last_result = copy.deepcopy(self.current_result)
                    #     self.time_last = time.time()
                    self.img_hk_draw_finish = False
                    # ret_val, self.img = camera.read()
                    result = inference_detector(self.model, self.img_hk)
                    pred_instances = result.pred_instances
                    pred_instances = pred_instances[pred_instances.scores > self.args.score_thr]
                    # plate_instances = pred_instances[pred_instances.labels==0]
                    # print(plate_instances)
                    labels,scores,bboxes = pred_instances.labels.cpu().numpy(),pred_instances.scores.cpu().numpy(),pred_instances.bboxes.cpu().numpy()
                except Exception as e:
                    print(str(e))
                    self.current_result_hk={"labels":[],"scores":[],"bbox":[]}
                    self.ready_hk = False
                else:
                    #TODO 分割
                    filter_edge = np.array([],dtype=np.bool_)
                    for b in bboxes:
                        filter_edge = np.append(filter_edge,np.array_equal(self.edge_flag,(b-self.edge_hk)>0))
                    labels,scores,bboxes = labels[filter_edge],scores[filter_edge],bboxes[filter_edge]
                    # labels.astype("str")
                    # labels,bboxes = self.change_label(labels,bboxes)
                    # print("22",labels,scores,bboxes)
                    # self.current_result["bbox"] = [self.y_pixel2world(point) if i%2 else self.x_pixel2world(point) for i,point in enumerate(bboxes.tolist())]
                    # bboxes_copy = bboxes.copy()
                    bboxes_copy = bboxes.flatten().reshape(bboxes.shape)
                    bboxes_copy[:, [ 0 , 2 ]] = bboxes_copy[:, [ 2 , 0 ]]#交换两列
                    self.current_result_hk["bbox"] = np.concatenate((bboxes,bboxes_copy),axis=1)#左上，右下，右上，左下
                    self.current_result_hk["labels"] = labels
                    self.current_result_hk["scores"] = scores

                    # mapping depth values
                    depth_maps = []
                    for bbox in bboxes:
                        x1, y1, x2, y2 = map(int, bbox)
                        x1, y1 = max(0, x1), max(0, y1)
                        x2, y2 = min(self.depth_img.shape[1], x2), min(self.depth_img.shape[0], y2)
                        depth_region = self.depth_img[y1:y2, x1:x2]
                        depth_maps.append(depth_region)
                    self.current_result_hk["depth_maps"] = depth_maps

                    self.ready_hk = True
                    # self.check_daocai_ready()
                    # print("daocaiready",self.daocai_ready)
                    indices = [i for i in range(len(bboxes))]
                    for index, bbox, label_id in zip(indices, bboxes, labels):
                        cv2.rectangle(self.img_hk, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (0, 255, 0),1)
                        cv2.putText(self.img_hk,str(label_id),(int(bbox[0]), int(bbox[1])),cv2.FONT_HERSHEY_SIMPLEX,0.6, (0, 255, 0), 1)
                    self.img_draw_finish_hk=True
                    # cv2.namedWindow('camera',cv2.WINDOW_NORMAL)
                    # cv2.imshow('camera', self.img)
                    # ch = cv2.waitKey(1)
                    # if ch == 27 or ch == ord('q') or ch == ord('Q'):
                    #     break
            else:
                self.hk_reader.destroy_camera()
                print("no data[0x%x]" % ret)
    # def process(self):
    #     while True:
    #         time.sleep(0.01)
    #         if self.ready and self.handrobot.ready:#机器人就绪:
    #             http_order.upload_plate_state(self.current_result)
    #             if bool(http_order.current_order):
    #                 format_current_order = http_order.format_current_order()#菜品类别和数量
    #                 for index,order in format_current_order.items():#处理订单
    #                     in_plate ,out_plate= list(order.keys())[0],list(order.values())[0]
    #                     for bbox, label_id in zip(self.current_result["bbox"], self.current_result["labels"]):
    #                         # if int(label_id)==out_plate:#餐格
    #                         if not self.flow[0]:
    #                             if self.handrobot.robot_do(in_plate,msg='取餐'):
    #                                 self.flow[0],self.flow[1]= True,True
    #                     for bbox, label_id in zip(self.current_result["bbox"], self.current_result["labels"]):
    #                         if int(label_id)==out_plate:#餐格
    #                             if  self.flow[1]:
    #                                 out_plate_center = (int((bbox[0]+bbox[2])/2),int((bbox[1]+bbox[3])/2))
    #                                 world_coord = self.pixel2world(out_plate_center)
    #                                 print(int(label_id),world_coord)
    #                                 if self.handrobot.robot_do_coord(world_coord):
    #                                     self.flow[2]= True

    #                     if self.flow[2]:
    #                         if self.handrobot.daocai(self.daocai_ready):
    #                             self.flow = [False,False,False]
    #                     break
    #         else:
    #             print("视觉未开启,机器人未就绪")

if __name__ == "__main__":
    
    stereovision = StereoVision()

    stereovision.return_vision_result()
